package com.innodealing.kafka.config;

import org.apache.kafka.clients.consumer.internals.AbstractPartitionAssignor;
import org.apache.kafka.common.TopicPartition;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 自定义Hash分配
 * @author penghai
 */
public class CustomHashModAssignor extends AbstractPartitionAssignor {

    private static final int INITIAL_CAPACITY = 64;

    @Override
    public String name() {
        return "custom-hash-mod";
    }

    @Override
    public Map<String, List<TopicPartition>> assign(Map<String, Integer> partitionsPerTopic,
                                                    Map<String, Subscription> subscriptions) {
        Map<String, List<TopicPartition>> assignment = new HashMap<>(INITIAL_CAPACITY);
        List<String> allConsumers = new ArrayList<>(subscriptions.keySet());
        // 保证排序，以使分配尽可能稳定
        allConsumers.sort(String::compareTo);

        int numConsumers = allConsumers.size();

        for (Map.Entry<String, Integer> entry : partitionsPerTopic.entrySet()) {
            String topic = entry.getKey();
            int numPartitions = entry.getValue();

            // 每个 Topic 都只有一个分区，这里假定 numPartitions = 1
            if (numPartitions != 1) {
                continue;
            }

            TopicPartition partition = new TopicPartition(topic, 0);
            int hashCode = topic.hashCode();
            int consumerIndex = Math.abs(hashCode) % numConsumers;
            String selectedConsumer = allConsumers.get(consumerIndex);

            assignment.computeIfAbsent(selectedConsumer, k -> new ArrayList<>()).add(partition);
        }

        // 确保所有消费者都在分配映射中，即使它们没有分配到分区
        for (String consumer : allConsumers) {
            assignment.computeIfAbsent(consumer, k -> new ArrayList<>());
        }

        return assignment;
    }
}
